/* 
 *  Copyright (c) 2020-2021 Xuhpclab. All rights reserved.
 *  Licensed under the MIT License.
 *  See LICENSE file for more information.
 */

#include <map>
#include <string>
#include <sys/stat.h>
#include <math.h>

#include "dr_api.h"
#include "drmgr.h"
#include "drcctlib.h"

using namespace std;
//extern int main_flag;

#define DRCCTLIB_PRINTF(_FORMAT, _ARGS...) \
    DRCCTLIB_PRINTF_TEMPLATE("thread_bottleneck_detection", _FORMAT, ##_ARGS)
#define DRCCTLIB_EXIT_PROCESS(_FORMAT, _ARGS...) \
    DRCCTLIB_CLIENT_EXIT_PROCESS_TEMPLATE("thread_bottleneck_detection", _FORMAT, ##_ARGS)

#define MAX_CLIENT_CCT_PRINT_DEPTH 10

#define SAMPLE_RUN
#ifdef SAMPLE_RUN
#    define UNITE_NUM 100000000
#    define SAMPLE_NUM 10000000
#endif

static std::string g_folder_name;
static int tls_idx;

typedef struct _cache_contention_t {
    int miss_count;
    int inval_count;
} cache_contention_t;

typedef struct _per_thread_t {
    long long int thread_local_lock_contention_count;
    long long int thread_local_lock_visit_count;
    int pt_id;
#ifdef SAMPLE_RUN
    uint32_t cur_mem_idx;
#endif
    file_t output_file;
// #define DEBUG_THREAD_BOTTLENECK
#ifdef DEBUG_THREAD_BOTTLENECK
    file_t log_file;
#endif
} per_thread_t;

//map<int, cache_contention_t> *thread_count_map = new map<int, cache_contention_t>();
static long long int total_lock_contention_count = 0;
static long long int total_lock_visit_count = 0;

inline bool
ContentioinTest(uint32_t bitmap, int pt_id) {
    if ((bitmap & (~(1 << pt_id))) == 0) {
        return false; //the same thread, do not happen contention
    } else {
        return true; //happen contention
    }
}

// client want to do
inline void
DoWhatClientWantTodo(void *drcontext, per_thread_t *pt, 
                        context_handle_t cur_ctxt_hndl, app_pc cur_addr)
{
    lock_handle_t *lock_hndl = 
                    get_lock_shadow_memory_space((void *)cur_addr);
    
    if(lock_hndl == NULL) {
        //this memory access is not a lock
        
    } else if (lock_hndl->lock_init == 1) {
	DRCCTLIB_PRINTF("lock visit addr : %x\n", (uint64_t)cur_addr);
        pt->thread_local_lock_visit_count++;
        uint32_t bitmap = lock_hndl->lock_owner_bitmap;
        if (ContentioinTest(bitmap, pt->pt_id)) {
		DRCCTLIB_PRINTF("lock contention addr : %x\n", (uint64_t)cur_addr);
            pt->thread_local_lock_contention_count++;//后续可以使用二元组将发生的冲突数与引发这些冲突的锁变量绑定
        }
    }
}

static inline void
InstrumentPerBBCache(void *drcontext, context_handle_t ctxt_hndl, int32_t slot_num,
                     int32_t mem_ref_num, mem_ref_msg_t *mem_ref_start, void **data)
{
    per_thread_t *pt;
    if (*data != NULL) {
        pt = (per_thread_t *)*data;
    } else {
        pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
        *data = pt;
    }
    //DRCCTLIB_PRINTF("client addr : %x\n", &main_flag);
    //printf("client: main flag is %d\n", read_main_flag());
    
    for (int32_t i = 0; i < mem_ref_num; i++) {
        if (mem_ref_start[i].slot >= slot_num) {
            break;
        }
        #ifdef SAMPLE_RUN
        pt->cur_mem_idx++;
        if (pt->cur_mem_idx > 0 && (pt->cur_mem_idx % UNITE_NUM <= SAMPLE_NUM) && read_main_flag()) {
        //dr_printf("id is %d : cur_mem_idx is %d\n", drcctlib_get_thread_id(), pt->cur_mem_idx);
          //printf("do client\n");
          DoWhatClientWantTodo(drcontext, pt, ctxt_hndl + mem_ref_start[i].slot,
                             mem_ref_start[i].addr);
        } else {
          //dr_printf("id is %d : cur_mem_idx is %d \n", drcctlib_get_thread_id(), pt->cur_mem_idx);
        }
        #else
        DoWhatClientWantTodo(drcontext, pt, ctxt_hndl + mem_ref_start[i].slot,
                             mem_ref_start[i].addr);
        #endif
    }
}


#ifdef DEBUG_CACHE_CONTENTION
static void
ThreadDebugFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.debug.log", g_folder_name.c_str(), id);
    pt->log_file = dr_open_file(name, DR_FILE_WRITE_APPEND | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->log_file != INVALID_FILE);
}
#endif

static void
ThreadOutputFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.topn.log", g_folder_name.c_str(), id);
    pt->output_file = dr_open_file(name, DR_FILE_WRITE_OVERWRITE | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->output_file != INVALID_FILE);
}

static void
ClientThreadStart(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)dr_thread_alloc(drcontext, sizeof(per_thread_t));
    if (pt == NULL) {
        DRCCTLIB_EXIT_PROCESS("pt == NULL");
    }
    drmgr_set_tls_field(drcontext, tls_idx, (void *)pt);

    ThreadOutputFileInit(pt);
    pt->pt_id = drcctlib_get_thread_id();
    pt->thread_local_lock_contention_count = 0;
    pt->thread_local_lock_visit_count = 0;
    #ifdef SAMPLE_RUN
    pt->cur_mem_idx=0;
    #endif
#ifdef DEBUG_CACHE_CONTENTION
    ThreadDebugFileInit(pt);
#endif
}

void
PrintResult(per_thread_t *pt) 
{
    int cur_thread_id = pt->pt_id;//drcctlib_get_thread_id();
    dr_fprintf(pt->output_file, "current thread id is %d\n ===========================================\n\n", cur_thread_id);
    /*map<int, cache_contention_t>::iterator it = (*thread_count_map).find(cur_thread_id);
    if (it == (*thread_count_map).end()) {
        dr_fprintf(pt->output_file, "current thread doesn't exist! Something goes wrong!!\n");
        return ;
    }*/
    dr_fprintf(pt->output_file, "total lock access nums are %lld  total lock contention are %lld \n\n",
										total_lock_visit_count, total_lock_contention_count);
    dr_fprintf(pt->output_file, "total lock contention / total lock access percentage : %2f% \n\n", 
										(double)total_lock_contention_count / (double)total_lock_visit_count * 100);
    //dr_fprintf(pt->output_file, "total cache inval / total mem ref percentage : %2f% \n\n",
										//(double)total_invalidation_count / (double)total_mem_ref_count * 100);
    
    //dr_fprintf(pt->output_file, "total cold miss nums are : %lld\n", total_cold_miss_count);
    dr_fprintf(pt->output_file, "======================================\n\n");
    dr_fprintf(pt->output_file, "thread lock contention count : %lld  total lock contention count : %lld  cur_thread contention / total contention percentage : %2f% \n\n", 
                                                                                pt->thread_local_lock_contention_count, total_lock_visit_count, 
                                                                                    ((double)pt->thread_local_lock_contention_count / total_lock_visit_count) * 100);
    /*dr_fprintf(pt->output_file, "thread inval count : %d  total inval count : %lld cur_thread inval / total inval  percentage : %2f% \n\n", 
                                                                                pt->thread_local_invalidation_count, total_invalidation_count,
                                                                                    ((double)pt->thread_local_invalidation_count / total_invalidation_count) * 100);
    */
}

static void
ClientThreadEnd(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
    total_lock_contention_count += pt->thread_local_lock_contention_count;
    total_lock_visit_count += pt->thread_local_lock_visit_count;
    //dr_printf("pt->id is %d\n", pt->pt_id);
    //dr_printf("drcctlib_get_thread_id is %d\n", drcctlib_get_thread_id());
    PrintResult(pt);
    dr_close_file(pt->output_file);
#ifdef DEBUG_CACHE_CONTENTION
    dr_close_file(pt->log_file);
#endif
    dr_thread_free(drcontext, pt, sizeof(per_thread_t));
}

static void
ClientInit(int argc, const char *argv[])
{
    char name[MAXIMUM_FILEPATH] = "";
    DRCCTLIB_INIT_LOG_FILE_NAME(
        name, "drcctlib_thread_bottleneck_detection", "out");
    g_folder_name.assign(name, strlen(name));
    mkdir(g_folder_name.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
}

static void
ClientExit(void)
{
    drcctlib_exit();
    if (!drmgr_unregister_thread_init_event(ClientThreadStart) ||
        !drmgr_unregister_thread_exit_event(ClientThreadEnd) ||
        !drmgr_unregister_tls_field(tls_idx)) {
        DRCCTLIB_PRINTF(
            "ERROR: drcctlib_thread_bottleneck_detection failed to unregister in ClientExit");
    }
    drmgr_exit();
}

#ifdef __cplusplus
extern "C" {
#endif

DR_EXPORT void
dr_client_main(client_id_t id, int argc, const char *argv[])
{
    dr_set_client_name("DynamoRIO Client 'drcctlib_thread_bottleneck_detection'",
                       "http://dynamorio.org/issues");
    ClientInit(argc, argv);

    if (!drmgr_init()) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_thread_bottleneck_detection unable to initialize drmgr");
    }
    drmgr_priority_t thread_init_pri = { sizeof(thread_init_pri),
                                         "drcctlib_thread_bottleneck-thread_init", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_priority_t thread_exit_pri = { sizeof(thread_exit_pri),
                                         "drcctlib_thread_bottleneck-thread-exit", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_register_thread_init_event_ex(ClientThreadStart, &thread_init_pri);
    drmgr_register_thread_exit_event_ex(ClientThreadEnd, &thread_exit_pri);
    tls_idx = drmgr_register_tls_field();
    if (tls_idx == -1) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_thread_bottleneck_detection drmgr_register_tls_field fail");
    }
    drcctlib_init_ex(DRCCTLIB_FILTER_MEM_ACCESS_INSTR, INVALID_FILE, NULL, NULL,
                     InstrumentPerBBCache,
                     DRCCTLIB_COLLECT_DATA_CENTRIC_MESSAGE | DRCCTLIB_CACHE_MODE |
                         DRCCTLIB_CACHE_MEMEORY_ACCESS_ADDR);
    dr_register_exit_event(ClientExit);
}

#ifdef __cplusplus
}
#endif
